#ifndef TYPE_HPP
#define TYPE_HPP

#include <string>
#include <map>
#include <vector>
using namespace std;

enum AllTypes {
    TYPE_INT,
    TYPE_BOOL,
    TYPE_STRING,
    TYPE_CHAR,
    TYPE_DOUBLE,

    // seperator: types before this are primitives.
    TYPE_PRIMITIVE, 
    TYPE_DEPENDENT,

    TYPE_ARROW,
    TYPE_LIST,
    TYPE_TUPLE,
    TYPE_RECORD,
    TYPE_TRAIT,

    TYPE_POLY,
    TYPE_CONSTRUCTOR,

    TYPE_FUNC,
    TYPE_APPLY,
    TYPE_VAR,

    TYPE_ANY,
    TYPE_GENERIC_VAR,
    TYPE_GENERIC_TRAIT,
};

enum OperatorType {
    OPERATOR_ADD,
    OPERATOR_SUB,
    OPERATOR_MUL,
    OPERATOR_DIV,
    OPERATOR_MOD,
    OPERATOR_SHIFT,
};

class TypeTable;

class Type {
protected:
    AllTypes _type_code;
    bool     _is_substing;

public:
    Type(AllTypes t) : _type_code(t) {
        _is_substing = false;
    }

    virtual Type* expand(TypeTable* type_table) { return this;   }
    virtual Type* substitute(TypeTable* type_table) { return this; }

    virtual ~Type() {}

    virtual bool support_operator(OperatorType ot, Type* other) {
        return false;
    }

    bool is_primitive() {
        return _type_code < TYPE_PRIMITIVE;
    }

    bool is_arrow() {
        return _type_code == TYPE_ARROW;
    }

    bool is_list() {
        return _type_code == TYPE_LIST;
    }

    bool is_record() {
        return _type_code == TYPE_RECORD;
    }

    bool is_trait() {
        return _type_code == TYPE_TRAIT;
    }

    bool is_tuple() {
        return _type_code == TYPE_TUPLE;
    }

    bool is_type_con() {
        return _type_code == TYPE_CONSTRUCTOR;
    }

    bool is_type_func() {
        return _type_code == TYPE_FUNC;
    }

    bool is_type_var() {
        return _type_code == TYPE_VAR;
    }

    bool is_type_apply() {
        return _type_code == TYPE_APPLY;
    }

    bool is_poly() {
        return _type_code == TYPE_POLY;
    }

    virtual string to_string() {
        return "";
    }

    virtual bool equals(Type* t) {
        return false;
    }

    bool is_substing() {
        return _is_substing;
    }

    void set_substing(bool is_substing) {
        _is_substing = is_substing;
    }
};

class PrimitiveType : public Type {
public:
    PrimitiveType(AllTypes t) : Type(t) {}

    virtual string to_string() {
        return "";
    }

    virtual bool equals(Type* t) {
        return this == t;
    }
};

class IntType : public PrimitiveType {
private:
    IntType() : PrimitiveType(TYPE_INT) {}

public:
    static IntType* get_instance() {
        static IntType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Int";
    }

    virtual bool support_operator(OperatorType ot, Type* other);
};

class BoolType : public PrimitiveType {
private:
    BoolType() : PrimitiveType(TYPE_BOOL) {}

public:
    static BoolType* get_instance() {
        static BoolType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Bool";
    }
};

class DoubleType : public PrimitiveType {
private:
    DoubleType(): PrimitiveType(TYPE_DOUBLE) {}

public:
    static DoubleType* get_instance() {
        static DoubleType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Double";
    }

    virtual bool support_operator(OperatorType ot, Type* other);
};

class UnitType : public PrimitiveType {
private:
    UnitType() : PrimitiveType(TYPE_BOOL) {}

public:
    static UnitType* get_instance() {
        static UnitType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Unit";
    }
};

class StringType : public PrimitiveType {
private:
    StringType() : PrimitiveType(TYPE_STRING) { }

public:
    static StringType* get_instance() {
        static StringType instance;
        return &instance;
    }

    virtual string to_string() {
        return "String";
    }

    virtual bool support_operator(OperatorType ot, Type* other);
};

class CharType : public PrimitiveType {
private:
    CharType() : PrimitiveType(TYPE_CHAR) {}

public:
    static CharType* get_instance() {
        static CharType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Char";
    }
    virtual bool support_operator(OperatorType ot, Type* other);
};

// Type for *->*
class ArrowType : public Type {
private:
    static const int LEN = 128;
    Type* _src;
    Type* _dst;

public:
    ArrowType(Type* src, Type* dst) : Type(TYPE_ARROW),
        _src(src), _dst(dst) {};
    ~ArrowType();

    Type* src() { return _src; }
    Type* dst() { return _dst; }

    void set_src(Type* t) { _src = t; }
    void set_dst(Type* t) { _dst = t; }

    virtual Type* substitute(TypeTable* type_table);
    virtual Type* expand(TypeTable* type_table) { return nullptr; }

    virtual string to_string();

    virtual bool equals(Type* t);
    Type* try_apply(Type* left);
};

class Token;

class UserDefType : public Type {
private:
    char*  _name;

public:
    UserDefType(Token* t) : Type(TYPE_CONSTRUCTOR) {};
    ~UserDefType();

    const char* name() { return _name; }
};

class AnyType : public Type {
private:
    AnyType() : Type(TYPE_ANY) {}

public:
    static AnyType* get_instance() {
        static AnyType instance;
        return &instance;
    }

    virtual string to_string() {
        return "Any";
    }

    virtual bool equals(Type* t) {
        return this == t;
    }
};

class Node;
class TypeFunction : public Type {
private:
    vector<Type*>  _args;
    Type*          _body;
    string         _name;

public:
    TypeFunction(Type* body) : Type(TYPE_FUNC),
        _body(body) {}

    vector<Type*>& args() { return _args; }
    Type* body()     { return _body; }
    string& name()   { return _name; }

    void set_body(Type* body) { _body = body; }
    void set_name(string name) { _name = name; }

    virtual Type* expand(TypeTable* type_table);
    virtual Type* expand_record(TypeTable* type_table);
    virtual Type* expand_trait(TypeTable* type_table);
};

class TypeVar : public Type {
private:
    string _name;
    TypeTable* _env;

public:
    TypeVar(string name, TypeTable* env) : Type(TYPE_VAR), _name(name), _env(env) {
    }

    string& name()     { return _name; }
    TypeTable* env()    { return _env; }

    virtual string to_string() {
        return "TypeVar(" + _name + ")";
    }

    virtual Type* get_type(TypeTable* type_table);
    virtual Type* substitute(TypeTable* type_table);
    virtual Type* expand(TypeTable* type_table) { return nullptr; }
    
    // 新增静态方法：每次调用创建一个新的TypeVarNode，名称为"?Tx"，x自增
    static TypeVar* make_type_var(TypeTable* env) {
        static int counter = 0;
        string name = "?T" + std::to_string(counter++);
        return new TypeVar(name, env);
    }

    virtual bool equals(Type* t);
};

class ListType : public Type {
private:
    Type*   _elm_type;

public:

    ListType(Type* ety) : Type(TYPE_LIST),
        _elm_type(ety) {}

    Type* elm_type()    { return _elm_type; }

    virtual bool equals(Type* t);
    virtual string to_string() {
        return "List(" + _elm_type->to_string() + ")";
    }

    virtual bool support_operator(OperatorType ot, Type* other);
    virtual Type* substitute(TypeTable* type_table);
};

class TupleType : public Type {
private:
    vector<Type*> _types;

public:
    TupleType(vector<Type*> types) : Type(TYPE_TUPLE), _types(types) {}

    vector<Type*>& types() { return _types; }

    virtual Type* substitute(TypeTable* type_table);
    virtual bool equals(Type* t);
    virtual string to_string();
};

class RecordType : public Type {
private:
    string _name;
    map<string, Type*> _fields;
    vector<Type*> _field_types;

    map<string, Type*> _method_types;
    map<string, Type*> _generic_args;

public:
    RecordType(string name, map<string, Type*> fields, vector<Type*> field_types) : Type(TYPE_RECORD),
        _name(name), _fields(fields), _field_types(field_types) {}

    virtual bool support_operator(OperatorType ot, Type* other);

    string& name() { return _name; }
    map<string, Type*>& fields() { return _fields; }
    vector<Type*>& field_types() { return _field_types; }

    Type* get_field_type(string& name) {
        return _fields[name];
    }

    virtual Type* substitute(TypeTable* type_table);

    void set_methods_types(map<string, Type*> methods) {
        _method_types = methods;
    }
    
    void add_method_type(string name, Type* type) {
        _method_types[name] = type;
    }

    Type* get_method_type(string& name) {
        return _method_types[name];
    }

    bool has_method(string name) {
        return _method_types.find(name) != _method_types.end();
    }

    map<string, Type*>& generic_args() { return _generic_args; }

    void add_generic_arg(string name, Type* type) {
        _generic_args[name] = type;
    }

    Type* get_generic_arg(string name) {
        auto it = _generic_args.find(name);
        if (it != _generic_args.end()) {
            return it->second;
        }
        return nullptr;
    }
};

class TraitType : public Type {
private:
    string _name;
    map<string, Type*> _methods;

public:
    TraitType(string name, map<string, Type*> methods) : Type(TYPE_TRAIT),
        _name(name), _methods(methods) {}

    map<string, Type*>& methods() { return _methods; }
    bool has_method(string name) {
        return _methods.find(name) != _methods.end();
    }
    string name() { return _name; }
};

class PolyType : public Type {
public:
    Type* type_constructor;
    vector<Type*> type_args;

    PolyType(Type* tc, vector<Type*> ta) : Type(TYPE_POLY),
        type_constructor(tc), type_args(ta) {}

    virtual Type* expand(TypeTable* type_table);
};

class TypeApply : public Type {
public:
    Type* type_constructor;
    vector<Type*> type_args;

    TypeApply(Type* tc, vector<Type*> ta) : Type(TYPE_APPLY),
        type_constructor(tc), type_args(ta) {}

    virtual string to_string() {
        return "TypeApply";
    }

    virtual Type* expand(TypeTable* type_table);
};

class TypeDependentConst : public Type {
public:
    int     value;

    TypeDependentConst(int v) : Type(TYPE_DEPENDENT), value(v) {}
};

class TypeTable {
private:
    map<string, Type*>    _type_map;
    map<string, Type*>    _type_var_map;
    TypeTable*            _parent;

    // 泛型类型变量，按顺序声明，所以使用vector存储
    vector<string>          _generic_names;
    map<string, Type*>      _substitution_map;

public:
    TypeTable();
    TypeTable(TypeTable* parent);

    Type* get_expanded_type(string& name);
    map<string, Type*>& type_map()      { return _type_map; }
    TypeTable* parent()                   { return _parent; }

    void add(string name, Type* t)     { _type_map[name] = t; }
    void add_type_var(string& name, Type* t) { _type_var_map[name] = t; }
    void add_generic_name(Node* n);

    vector<string>& generic_names() { return _generic_names; }

    vector<Type*> get_generic_types() {
        vector<Type*> types;
        for (auto it = _generic_names.begin(); it != _generic_names.end(); it++) {
            types.push_back(_substitution_map[*it]);
        }
        return types;
    }

    void set_generic_name(int index, Type* t)   { 
        _substitution_map[_generic_names[index]] = t;
    }

    void add_substitution(string name, Type* t) { 
        _substitution_map[name] = t;
    }
    
    Type* get_substitution(string& name);

    void clear_instantiation();

    Type* ufset_find(Type* t);
    void ufset_union(Type* t1, Type* t2);
};

#endif

